import random
import warnings
import numpy as np
import pandas as pd
import argparse

from sklearn.datasets import load_boston, load_iris


def get_arguments():
    parser = argparse.ArgumentParser(description='MyDecisionTree')
    parser.add_argument('--dataset', type=int, default=1, choices=(1, 2),
                        help='the name of dataset,'
                             '1: the iris_dataset,'
                             '2: the boston dataset')
    parser.add_argument('--test_size', type=float, default=0.33, help='the proportion of test data')
    parser.add_argument('--random_state', type=int, default=42, help='the random seed of dataset split')

    parser.add_argument('--threshold')

    args = parser.parse_args()
    return args


class MyPreprocessing:
    def __init__(self, parser):
        self.dataset = parser.dataset
        self.random_state = parser.random_state
        self.test_size = parser.test_size

    def load_dataset(self):
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            if self.dataset == 1:
                dataset = load_iris()
            elif self.dataset == 2:
                dataset = load_boston()
            else:
                raise ValueError("Please choose right dataset~", self.dataset)
            datas = pd.DataFrame(dataset.data)
            datas.columns = dataset.feature_names
            target = dataset.target
            datas.insert(loc=len(datas.columns), column='target', value=target)
            return datas

    def split_dataset(self, datas):
        assert 0 < self.test_size < 1, "Please choose right test size between 0 and 1"
        test_num = int(self.test_size * datas.shape[0])
        np.random.seed(self.random_state)
        i = random.sample(range(0, datas.shape[0]), test_num)
        test = datas.loc[i]
        train = datas.drop(labels=i, axis=0)
        X_train = train.iloc[:, :-1]
        y_train = train.iloc[:, -1]
        X_test = test.iloc[:, :-1]
        y_test = test.iloc[:, -1]
        return X_train, y_train, X_test, y_test


class CART_decision_tree(object):
    """
    一个递归定义的数据结构，用于存储树。
    每个节点可以包含其他节点作为其子节点。
    """

    def __init__(self, tree='cls', criterion='gini', max_depth=None, TN=0, TG=0.0):

        self.feature = None  # 特征对应列号(索引)
        self.featurename = None
        self.label = None
        self.n_samples = None
        self.gain = None
        self.left = None
        self.right = None
        self.threshold = None
        self.depth = 0

        self.root = None
        self.criterion = criterion
        self.max_depth = max_depth
        self.TN = TN
        self.TG = TG
        self.tree = tree

    def fit(self, features, target):
        """
        该函数接受一个训练数据集，从中建立树状结构来做决定，或者做子节点来做进一步的查询。

        参数:
        features (dataframe): [n * p] n个p属性的观察数据样本
        target (series): [n] 目标值
        """
        self.root = CART_decision_tree()
        if self.tree == 'cls':
            self.root._grow_tree(features, target, self.criterion)
        else:
            self.root._grow_tree(features, target, 'mse')
        if self.max_depth != None:
            self.root._prune(self.max_depth)

    def predict(self, features):
        """
        该函数接受一个测试数据集，从中遍历CART树，输出一个 预测数组。

        参数:
        features (dataframe): [n * p] n个p属性的观察数据样本

        返回值:
        list: [n] 预测值
        """
        # 遍历每个测试实例，并反复调用内部方法"_predicts "进行预测
        # 如果当前节点属性低于或高于阈值，直到可以进行预测。
        predicts = []
        for col in range(features.shape[0]):
            test_feature = features.iloc[col]
            predict = self.root._predict(test_feature)
            predicts.append(predict)
        return predicts

    def print_tree(self):
        """
        用于打印决策树结构的辅助函数。
        """
        self.root._show_tree(0, '')

    def _grow_tree(self, features, target, criterion='gini'):
        """"
        内部函数通过将数据分割成两个分支来增长树的长度，基于一些阈值的值。目前的节点被分配到提供最佳分割的特征，并有相应的阈值。

        参数:
        features (dataframe): [n * p] n个p属性的观察数据样本
        target (series): [n] 目标值
        criterion (string): 分类标准。 default = 'gini'
        """
        self.n_samples = features.shape[0]
        # 如果所有的类都是相同的，节点被标记为该类并返回。终止检索。（待补充）
        if len(np.unique(target)) == 1:
            self.label = target.values[0]
            return
        if self.n_samples <= self.TN:
            return

        if criterion in {'gini', 'entropy'}:
            # 计算哪个类在数据中拥有最多的实例，并为节点分配该标签。（待补充）
            label = list(target)
            self.label = max(set(label), key=label.count)
        else:
            # 计算回归问题中的平均值，并为节点分配该标签。（待补充）
            self.label = np.mean(target)

        # 计算父节点的不纯度。（待补充）
        father_impurity = self._calc_impurity(criterion, target)

        best_gain = 0.0
        best_feature = None
        best_threshold = None
        # 迭代所有分类决策，以确定最佳信息增益、特征和阈值。
        features.insert(loc=len(features.columns), column='target', value=target)
        for col in range(features.shape[1]-1):
            # 找不同特征值
            feature = features[[features.columns[col]]]
            feature = np.unique(feature.values)
            feature.sort()

            # 按连续特征：单特征阈值法
            # 创建一个候选阈值的列表thresholds。(待补充）
            thresholds = []
            if len(feature) == 1:
                continue
            else:
                for k in range(len(feature) - 1):
                    thresholds.append((feature[k + 1] + feature[k])/2)

            # 遍历阈值
            for threshold in thresholds:
                # 计算左侧子节点的数据集标签。（待补充）
                left_target = features.loc[features[features.columns[col]] < threshold]['target']
                # 计算左侧子节点的不纯度标准。（调用self._calc_impurity方法）（待补充）
                left_impurity = self._calc_impurity(criterion, list(left_target))
                # 计算左侧子节点样本量/父节点样本量的比率。（待补充）
                left_ratio = len(left_target) / len(target)
                # 计算右侧子节点的数据集。（待补充）
                right_target = features.loc[features[features.columns[col]] > threshold]['target']
                # 计算右侧子节点的不纯度标准。（待补充）
                right_impurity = self._calc_impurity(criterion, list(right_target))
                # 计算右侧子节点样本量/父节点样本量的比率。（待补充）
                right_ratio = len(right_target) / len(target)
                # 计算信息增益。（待补充）
                decrease_impurity = father_impurity - (left_impurity * left_ratio + right_impurity * right_ratio)

                if decrease_impurity > best_gain:
                    best_gain = decrease_impurity
                    best_feature = col
                    best_threshold = threshold
        if best_feature <= self.TG:
            return
        self.gain = best_gain
        self.feature = best_feature
        self.featurename = features.columns[best_feature]
        self.threshold = best_threshold
        # 迭代创建子节点
        features = features.drop(['target'], axis=1)
        self._split_tree(features, target, criterion)

    def _split_tree(self, features, target, criterion):
        """"
        内部函数分割数据集，创建左、右子节点，并逐渐生长树。

        参数:
        features (dataframe): [n * p] n个p属性的观察数据样本
        target (series): [n] 目标值
        criterion (string): 分类标准。 default = 'gini'
        """
        features.insert(loc=len(features.columns), column='target', value=target)
        # 划分左侧子节点样本集合（左侧节点特征集合，左侧节点标签集合）
        left_target = features.loc[features[features.columns[self.feature]] < self.threshold]['target']
        left_feature = features.loc[features[features.columns[self.feature]] < self.threshold].iloc[:, :-1]

        # 递归树结构
        self.left = CART_decision_tree()
        self.left.depth = self.depth + 1
        self.left._grow_tree(left_feature, left_target, criterion)

        # 划分右侧子节点样本集合（右侧节点特征集合，右侧节点标签集合）
        right_target = features.loc[features[features.columns[self.feature]] > self.threshold]['target']
        right_feature = features.loc[features[features.columns[self.feature]] > self.threshold].iloc[:, :-1]

        self.right = CART_decision_tree()
        self.right.depth = self.depth + 1
        self.right._grow_tree(right_feature, right_target, criterion)

    def _calc_impurity(self, criterion, target):
        """"
            内部函数根据选择的措施计算杂质标准并返回数值。

            参数:
            criterion (string): 分类标准
            target (series): [n] 目标值

            返回值:
            float: 不纯度标准的值
            """
        length = len(target)
        labels = np.unique(target)
        prob = []
        target = list(target)
        for i, label in enumerate(labels):
            prob.append(target.count(label) / length)

        if criterion == 'gini':
            gini = 1 - np.power(prob, 2).sum()
            return gini
        elif criterion == 'mse':
            mse = np.power(target - np.mean(target), 2).sum() / length
            return mse
        else:
            # 计算信息增益
            if prob == 1 or prob == 0:
                entropy = 0
            else:
                entropy = -(prob * np.log2(prob)).sum()
            return entropy

    def _prune(self, max_depth):
        """"
            内部函数根据max_depth来调整树的深度。减少了对数据的过度拟合。
            参数:
            max_depth (int): 树的最大深度
            """
        if self.feature is None:
            return

        self.left._prune(max_depth)
        self.right._prune(max_depth)

        pruning = False

        if self.depth >= max_depth:
            pruning = True

        if pruning is True:
            self.left = None
            self.right = None
            self.feature = None

    def _predict(self, d):
        """"
            内部函数接收一行输入值并递归地检查一个阈值，直到返回一个 预测的标签被返回。

            参数:
            d (series): [1 * p] 观察到的p属性的数据样本

            返回值:
            string: 预测的标签
            """
        if self.feature != None:
            if d[self.feature] <= self.threshold:
                return self.left._predict(d)
            else:
                return self.right._predict(d)
        else:
            return self.label

    def _show_tree(self, depth, cond):
        """
            内部辅助函数，用于打印树形决策结构。
            """
        base = '|---' * depth + cond
        if self.feature != None:
            print(base + 'if ' + self.featurename + ' <= ' + f"{self.threshold:.2f}")
            self.left._show_tree(depth + 1, 'then ')
            self.right._show_tree(depth + 1, 'else ')
        else:
            print(base + '{class is: ' + str(self.label) + ', number of samples: ' + str(self.n_samples) + '}')

    def score(self, y_true, y_pred):
        """

        :param y_true: 真实值
        :param y_pred: 预测值
        :return: 分类指标/回归指标
        """
        if self.tree == 'cls':
            # 分类：交叉熵损失
            score = np.sum(y_true == y_pred) / len(y_true)
            # score = (y_true * np.log2(y_pred)).sum() / len(y_true)
        else:
            # 回归：均方误差
            score = np.power(y_true - y_pred, 2).sum() / len(y_true)
        return score


if __name__ == "__main__":
    # 加载参数
    parser = get_arguments()

    # boston房价回归树
    parser.dataset = 2
    # 1、导入数据集
    MyPreprocessing = MyPreprocessing(parser)
    datas = MyPreprocessing.load_dataset()
    # 2、数据集划分
    X_train, y_train, X_test, y_test = MyPreprocessing.split_dataset(datas)
    # 3、创建回归树实例
    regressionTree = CART_decision_tree(tree='mse', criterion='mse', max_depth=3)
    # 4、训练模型（调用fit)
    regressionTree.fit(X_train, y_train)
    # 5、输出树模型
    regressionTree.print_tree()
    # 6、预测，评分
    pred_train = regressionTree.predict(X_train)
    pred_test = regressionTree.predict(X_test)
    score_train = regressionTree.score(y_train, pred_train)
    score_test = regressionTree.score(y_test, pred_test)
    print("The score of training is {}".format(score_train))
    print("The score of testing is {}".format(score_test))

    # # iris分类树
    # parser.dataset = 1
    # # 1、导入数据集
    # MyPreprocessing = MyPreprocessing(parser)
    # datas = MyPreprocessing.load_dataset()
    # # 2、数据集划分
    # X_train, y_train, X_test, y_test = MyPreprocessing.split_dataset(datas)
    # # 3、创建回归树实例
    # classificationTree = CART_decision_tree(tree='cls', criterion='entropy', TN=1, TG=0.1)
    # # 4、训练模型（调用fit)
    # classificationTree.fit(X_train, y_train)
    # # 5、输出树模型
    # classificationTree.print_tree()
    # # 6、预测，评分
    # pred_train = classificationTree.predict(X_train)
    # pred_test = classificationTree.predict(X_test)
    # score_train = classificationTree.score(y_train, pred_train)
    # score_test = classificationTree.score(y_test, pred_test)
    # print("The score of training is {}".format(score_train))
    # print("The score of testing is {}".format(score_test))
